import os
import shutil
import cv2
import numpy as np
from pathlib import Path
import tkinter as tk
from tkinter import filedialog, messagebox, BooleanVar
import threading
from PIL import Image
import hashlib
from collections import defaultdict
import logging

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class ImageSelector:
    def __init__(self, process_images=False):
        self.supported_formats = {'.tiff', '.tif', '.png', '.jpg', '.jpeg'}
        self.max_similar_images = 8
        self.process_images = process_images
        
    def calculate_image_hash(self, image_path):
        """计算图片的感知哈希值"""
        try:
            # 读取图片
            img = cv2.imread(str(image_path))
            if img is None:
                return None
            
            # 转换为灰度图
            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            
            # 调整大小为8x8
            resized = cv2.resize(gray, (8, 8))
            
            # 计算平均值
            avg = resized.mean()
            
            # 生成哈希值
            hash_value = 0
            for i in range(8):
                for j in range(8):
                    if resized[i, j] > avg:
                        hash_value |= 1 << (i * 8 + j)
            
            return hash_value
        except Exception as e:
            logger.error(f"计算图片哈希值时出错 {image_path}: {e}")
            return None
    
    def hamming_distance(self, hash1, hash2):
        """计算两个哈希值的汉明距离"""
        if hash1 is None or hash2 is None:
            return float('inf')
        return bin(hash1 ^ hash2).count('1')
    
    def is_similar(self, hash1, hash2, threshold=10):
        """判断两个图片是否相似"""
        distance = self.hamming_distance(hash1, hash2)
        return distance <= threshold
    
    def get_image_files(self, folder_path):
        """获取文件夹中支持的图片文件"""
        image_files = []
        folder = Path(folder_path)
        
        if not folder.exists():
            return image_files
        
        for file_path in folder.iterdir():
            if file_path.is_file() and file_path.suffix.lower() in self.supported_formats:
                image_files.append(file_path)
        
        return sorted(image_files)
    
    def select_images(self, image_files):
        """筛选图片，每种相似图片最多保留8张"""
        if not image_files:
            return []
        
        # 预先计算所有图片的哈希值
        image_hashes = {}
        for image_path in image_files:
            image_hash = self.calculate_image_hash(image_path)
            if image_hash is not None:
                image_hashes[image_path] = image_hash
        
        selected_images = []
        processed_hashes = set()
        
        for image_path, image_hash in image_hashes.items():
            # 检查是否与已处理的图片相似
            is_similar_to_processed = False
            for processed_hash in processed_hashes:
                if self.is_similar(image_hash, processed_hash):
                    is_similar_to_processed = True
                    break
            
            if not is_similar_to_processed:
                # 找到所有与当前图片相似的图片
                similar_images = []
                for other_path, other_hash in image_hashes.items():
                    if self.is_similar(image_hash, other_hash):
                        similar_images.append(other_path)
                
                # 最多保留8张相似图片
                selected_images.extend(similar_images[:self.max_similar_images])
                processed_hashes.add(image_hash)
        
        return selected_images
    
    def process_folder(self, input_folder):
        """处理单个文件夹"""
        logger.info(f"处理文件夹: {input_folder}")
        
        # 获取图片文件
        image_files = self.get_image_files(input_folder)
        logger.info(f"找到 {len(image_files)} 个图片文件")
        
        if not image_files:
            logger.info(f"文件夹 {input_folder} 中没有找到支持的图片文件")
            return []
        
        # 筛选图片
        selected_images = self.select_images(image_files)
        logger.info(f"筛选后保留 {len(selected_images)} 张图片")
        
        return selected_images
    
    def create_output_folder(self, input_folder):
        """创建输出文件夹"""
        input_path = Path(input_folder)
        parent = input_path.parent
        folder_name = input_path.name
        output_folder = parent / f"{folder_name}_selected"
        
        # 如果输出文件夹已存在，先删除
        if output_folder.exists():
            shutil.rmtree(output_folder)
        
        # 创建输出文件夹
        output_folder.mkdir(parents=True, exist_ok=True)
        logger.info(f"创建输出文件夹: {output_folder}")
        
        return output_folder
    
    def process_image(self, image_path, index):
        """处理图片：根据索引进行不同的变换"""
        try:
            # 使用PIL读取图片以保持原始格式
            pil_img = Image.open(image_path)
            
            # 根据索引进行不同的处理
            cycle_index = index % 4
            
            if cycle_index == 0:
                # 第一张：保持不变
                processed_img = pil_img
                logger.info(f"处理图片: {image_path.name} - 保持不变")
            elif cycle_index == 1:
                # 第二张：水平翻转
                processed_img = pil_img.transpose(Image.FLIP_LEFT_RIGHT)
                logger.info(f"处理图片: {image_path.name} - 水平翻转")
            elif cycle_index == 2:
                # 第三张：垂直翻转
                processed_img = pil_img.transpose(Image.FLIP_TOP_BOTTOM)
                logger.info(f"处理图片: {image_path.name} - 垂直翻转")
            elif cycle_index == 3:
                # 第四张：旋转180度
                processed_img = pil_img.transpose(Image.ROTATE_180)
                logger.info(f"处理图片: {image_path.name} - 旋转180度")
            
            return processed_img
            
        except Exception as e:
            logger.error(f"处理图片时出错 {image_path}: {e}")
            return None
    
    def copy_selected_images(self, selected_images, output_folder):
        """复制选中的图片到输出文件夹"""
        for i, image_path in enumerate(selected_images):
            try:
                dest_path = output_folder / image_path.name
                
                if self.process_images:
                    # 处理图片
                    processed_img = self.process_image(image_path, i)
                    if processed_img is not None:
                        # 获取原始图片格式
                        original_format = None
                        try:
                            with Image.open(image_path) as original_img:
                                original_format = original_img.format
                        except:
                            pass
                        
                        # 使用原始格式保存
                        if original_format:
                            processed_img.save(dest_path, format=original_format)
                        else:
                            processed_img.save(dest_path)
                        logger.info(f"保存处理后的图片: {image_path.name} (格式: {original_format})")
                    else:
                        # 如果处理失败，复制原图
                        shutil.copy2(image_path, dest_path)
                        logger.info(f"复制原图: {image_path.name}")
                else:
                    # 直接复制原图
                    shutil.copy2(image_path, dest_path)
                    logger.info(f"复制图片: {image_path.name}")
                    
            except Exception as e:
                logger.error(f"复制图片失败 {image_path}: {e}")
    
    def process_directory_recursive(self, input_folder):
        """递归处理目录及其所有子目录"""
        input_path = Path(input_folder)
        
        if not input_path.exists():
            logger.error(f"输入路径不存在: {input_folder}")
            return
        
        # 创建主输出文件夹
        main_output_folder = self.create_output_folder(input_folder)
        
        # 递归处理所有文件夹
        self._process_folder_recursive(input_path, main_output_folder)
        
        logger.info(f"处理完成，输出文件夹: {main_output_folder}")
    
    def _process_folder_recursive(self, input_path, output_path):
        """递归处理文件夹及其所有子文件夹"""
        # 处理当前文件夹
        selected_images = self.process_folder(str(input_path))
        self.copy_selected_images(selected_images, output_path)
        
        # 递归处理所有子文件夹
        for item in input_path.iterdir():
            if item.is_dir():
                # 创建对应的子输出文件夹
                sub_output_folder = output_path / item.name
                sub_output_folder.mkdir(parents=True, exist_ok=True)
                
                # 递归处理子文件夹
                self._process_folder_recursive(item, sub_output_folder)

class ImageSelectorGUI:
    def __init__(self):
        self.selector = None
        self.setup_gui()
    
    def setup_gui(self):
        """设置GUI界面"""
        self.root = tk.Tk()
        self.root.title("图片筛选工具")
        self.root.geometry("600x800")
        self.root.resizable(True, True)
        
        # 设置中文字体
        try:
            self.root.option_add('*Font', ('Microsoft YaHei', 10))
        except:
            pass
        
        # 创建变量（必须在创建根窗口之后）
        self.process_images_var = BooleanVar()
        
        # 创建主框架
        main_frame = tk.Frame(self.root, padx=20, pady=20)
        main_frame.pack(fill=tk.BOTH, expand=True)
        
        # 标题
        title_label = tk.Label(main_frame, text="图片筛选工具", font=("Microsoft YaHei", 16, "bold"))
        title_label.pack(pady=(0, 20))
        
        # 说明文字
        description = tk.Label(main_frame, text="选择包含图片的文件夹，程序将自动筛选相似图片并保存", 
                              font=("Microsoft YaHei", 10))
        description.pack(pady=(0, 20))
        
        # 路径选择框架
        path_frame = tk.Frame(main_frame)
        path_frame.pack(fill=tk.X, pady=(0, 20))
        
        # 路径标签
        path_label = tk.Label(path_frame, text="输入文件夹:", font=("Microsoft YaHei", 10))
        path_label.pack(anchor=tk.W)
        
        # 路径输入框架
        input_frame = tk.Frame(path_frame)
        input_frame.pack(fill=tk.X, pady=(5, 0))
        
        self.path_var = tk.StringVar()
        self.path_entry = tk.Entry(input_frame, textvariable=self.path_var, font=("Microsoft YaHei", 10))
        self.path_entry.pack(side=tk.LEFT, fill=tk.X, expand=True)
        
        browse_button = tk.Button(input_frame, text="浏览", command=self.browse_folder, 
                                 font=("Microsoft YaHei", 10))
        browse_button.pack(side=tk.RIGHT, padx=(10, 0))
        
        # 图片处理选项
        process_frame = tk.Frame(main_frame)
        process_frame.pack(fill=tk.X, pady=(20, 0))
        
        process_checkbox = tk.Checkbutton(process_frame, text="对图片进行处理", 
                                         variable=self.process_images_var,
                                         font=("Microsoft YaHei", 10))
        process_checkbox.pack(anchor=tk.W)
        
        process_desc = tk.Label(process_frame, 
                               text="选中后，每组相似图片将按4张循环处理：第1张不变，第2张水平翻转，第3张垂直翻转，第4张旋转180度",
                               font=("Microsoft YaHei", 9), fg="gray", wraplength=500, justify=tk.LEFT)
        process_desc.pack(anchor=tk.W, pady=(5, 0))
        
        # 处理按钮
        self.process_button = tk.Button(main_frame, text="开始处理", command=self.start_processing,
                                       font=("Microsoft YaHei", 12, "bold"), bg="#4CAF50", fg="white",
                                       height=2)
        self.process_button.pack(fill=tk.X, pady=(20, 0))
        
        # 进度显示
        self.progress_label = tk.Label(main_frame, text="", font=("Microsoft YaHei", 10))
        self.progress_label.pack(pady=(10, 0))
        
        # 日志显示
        log_frame = tk.Frame(main_frame)
        log_frame.pack(fill=tk.BOTH, expand=True, pady=(20, 0))
        
        log_label = tk.Label(log_frame, text="处理日志:", font=("Microsoft YaHei", 10))
        log_label.pack(anchor=tk.W)
        
        # 创建文本框和滚动条
        text_frame = tk.Frame(log_frame)
        text_frame.pack(fill=tk.BOTH, expand=True, pady=(5, 0))
        
        self.log_text = tk.Text(text_frame, height=10, font=("Consolas", 9))
        scrollbar = tk.Scrollbar(text_frame, orient=tk.VERTICAL, command=self.log_text.yview)
        self.log_text.configure(yscrollcommand=scrollbar.set)
        
        self.log_text.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)
        scrollbar.pack(side=tk.RIGHT, fill=tk.Y)
        
        # 配置日志处理器
        self.setup_logging()
    
    def setup_logging(self):
        """设置日志显示"""
        class TextHandler(logging.Handler):
            def __init__(self, text_widget):
                logging.Handler.__init__(self)
                self.text_widget = text_widget
            
            def emit(self, record):
                msg = self.format(record)
                def append():
                    self.text_widget.insert(tk.END, msg + '\n')
                    self.text_widget.see(tk.END)
                self.text_widget.after(0, append)
        
        # 添加文本处理器
        text_handler = TextHandler(self.log_text)
        text_handler.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s'))
        logger.addHandler(text_handler)
    
    def browse_folder(self):
        """浏览文件夹"""
        folder_path = filedialog.askdirectory(title="选择包含图片的文件夹")
        if folder_path:
            self.path_var.set(folder_path)
    
    def start_processing(self):
        """开始处理"""
        input_folder = self.path_var.get().strip()
        
        if not input_folder:
            messagebox.showerror("错误", "请选择输入文件夹")
            return
        
        if not os.path.exists(input_folder):
            messagebox.showerror("错误", "选择的文件夹不存在")
            return
        
        # 禁用按钮
        self.process_button.config(state=tk.DISABLED, text="处理中...")
        self.progress_label.config(text="正在处理，请稍候...")
        
        # 在新线程中处理
        thread = threading.Thread(target=self.process_in_thread, args=(input_folder,))
        thread.daemon = True
        thread.start()
    
    def process_in_thread(self, input_folder):
        """在线程中处理"""
        try:
            # 根据用户选择创建图片筛选器
            self.selector = ImageSelector(process_images=self.process_images_var.get())
            self.selector.process_directory_recursive(input_folder)
            
            # 在主线程中更新UI
            self.root.after(0, self.processing_complete, True)
        except Exception as e:
            logger.error(f"处理过程中出错: {e}")
            self.root.after(0, self.processing_complete, False, str(e))
    
    def processing_complete(self, success, error_msg=None):
        """处理完成"""
        self.process_button.config(state=tk.NORMAL, text="开始处理")
        self.progress_label.config(text="")
        
        if success:
            messagebox.showinfo("完成", "图片筛选处理完成！")
        else:
            messagebox.showerror("错误", f"处理过程中出错：{error_msg}")
    
    def run(self):
        """运行GUI"""
        self.root.mainloop()

def main():
    """主函数"""
    try:
        app = ImageSelectorGUI()
        app.run()
    except Exception as e:
        print(f"程序启动失败: {e}")
        input("按回车键退出...")

if __name__ == "__main__":
    main() 